# client2.py (接收端)
import socket
import numpy as np
import threading
import socket
import cv2
import struct
import time

from lerobot.common.robot_devices.control_configs import ControlPipelineConfig
from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot
import draccus
import torch


cli_args = ['--robot.type=so101F', '--robot.cameras={}', '--control.type=teleoperate']

cfg = draccus.parse(config_class=ControlPipelineConfig, config_path=None, args=cli_args)
robot = ManipulatorRobot(cfg.robot)


if not robot.is_connected:
    robot.connect()

HOST = '47.109.17.68'
PORT = 8869

#HOST = '0.0.0.0'
#PORT = 65432

HEADER_SIZE = 4
DATA_SHAPE = (6,)

def connect_to_server():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((HOST, PORT))
    s.sendall(b'2')  # 标识为客户端2
    return s

def main():
    sock = connect_to_server()
    print("Receiver connected to server")
    buffer = bytearray()
    
    try:
        while True:
            # 高效接收数据
            chunk = sock.recv(4096)
            if not chunk: break
            buffer.extend(chunk)
            
            # 处理完整数据包
            while len(buffer) >= HEADER_SIZE:
                # 解析包头
                header = bytes(buffer[:HEADER_SIZE])
                data_size = int.from_bytes(header, 'big')
                
                # 检查是否收到完整数据
                if len(buffer) < HEADER_SIZE + data_size:
                    break
                
                # 提取数据包
                raw_data = buffer[HEADER_SIZE:HEADER_SIZE+data_size]
                del buffer[:HEADER_SIZE+data_size]
                
                # 反序列化numpy数组
                data = np.frombuffer(raw_data, dtype=np.float32)
                data = torch.from_numpy(data)
                print(f"Received: {data}")
                robot.send_action(data)

    except KeyboardInterrupt:
        print("\nReceiver stopped")
    finally:
        sock.close()


def video():
    # 连接配置
    SERVER_IP = '47.109.17.68'  # 服务端IP
    SERVER_PORT = 8870

    # 初始化摄像头
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    cap.set(cv2.CAP_PROP_FPS, 30)  # 设置帧率

    # 连接服务器
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect((SERVER_IP, SERVER_PORT))

    # 压缩参数
    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 80]  # 80%质量

    try:
        while cap.isOpened():
            # 读取帧
            ret, frame = cap.read()
            if not ret:
                break
            
            # 压缩图像
            _, buffer = cv2.imencode('.jpg', frame, encode_param)
            
            # 发送帧大小信息
            frame_size = len(buffer)
            client_socket.sendall(struct.pack('!I', frame_size))
            
            # 发送帧数据
            client_socket.sendall(buffer)
            
            # 控制传输速率 (约30fps)
            time.sleep(0.03)
            
    except KeyboardInterrupt:
        print("客户端关闭")
    except Exception as e:
        print(f"错误: {e}")
    finally:
        cap.release()
        client_socket.close()
if __name__ == "__main__":
    main_thread = threading.Thread(target=main, daemon=True)
    main_thread.start()
    
    v_thread = threading.Thread(target=video, daemon=True)
    v_thread.start()

    while True:
        time.sleep(1)

